package ftp

import (
	"fmt"
	"github.com/panjf2000/ants/v2"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

func Download(host string, port string, user string, pwd string, remotePath string, localPath string, poolSize int) {
	log.Printf("目标主机:%v,端口:%v\n用户名:%v,密码:***\n,远程目录:%v\n本地目录:%v\n",
		host, port, user, remotePath, localPath)
	sshClient, sftpClient := CreateClient(host, port, user, pwd)
	defer sshClient.Close()
	defer sftpClient.Close()

	//获取当前目录
	cwd, err := sftpClient.Getwd()
	if err != nil {
		log.Fatalln(err.Error())
	}
	log.Println("当前目录：", cwd)

	//显示文件/目录详情
	//fi, err := sftpClient.Lstat(cwd)
	//log.Println(fi)
	{
		//下载文件
		//将远程服务器的/bin/bash文件下载到本地
		remoteFileInfo, err := sftpClient.Stat(remotePath)
		if err != nil {
			panic(err)
		}
		//判断目录为空
		if remoteFileInfo.IsDir() {
			dir, err := sftpClient.ReadDir(remotePath)
			if err != nil {
				panic(err.Error())
			}
			if len(dir) == 0 {
				panic(fmt.Sprintf("远程目录:%v 为空", remotePath))
			}
		}
		files := GetFiles(sftpClient, localPath, remotePath)
		if pool, err := ants.NewPool(poolSize); err == nil {
			defer pool.Release()
			var wg sync.WaitGroup
			for i := range files {
				wg.Add(1)
				f := files[i]
				_ = pool.Submit(func() {
					downLoadSingleFile(sftpClient, f)
					wg.Done()
				})
			}
			wg.Wait()
		}

		/*
			wg.Add(1)
			go downLoadFile(
				&wg,
				sftpClient,
				remoteFileInfo,
				localPath,
				remotePath,
			)*/
		//递归下载

	}
}

func CreateClient(host string, port string, user string, pwd string) (sshClient *ssh.Client, sftpClient *sftp.Client) {
	sshConfig := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pwd),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		ClientVersion:   "",
		Timeout:         10 * time.Second,
	}

	//建立与SSH服务器的连接
	sshClient, err := ssh.Dial("tcp", host+":"+port, sshConfig)
	if err != nil {
		panic(err)
	}

	sftpClient, err = sftp.NewClient(sshClient)
	if err != nil {
		panic(err)
	}
	return sshClient, sftpClient
}

func downLoadFile(wg *sync.WaitGroup, sftpClient *sftp.Client, remoteFileInfo os.FileInfo, localPath string, remotePath string) {
	defer wg.Done()
	remoteFileName := remoteFileInfo.Name()
	fullLocalPath := localPath + "/" + remoteFileName
	if remoteFileInfo.IsDir() {
		dir, err := sftpClient.ReadDir(remotePath)
		if err != nil {
			panic(err.Error())
		}
		for _, f := range dir {
			fullRemotePath := remotePath + "/" + f.Name()
			if f.IsDir() {
				localPath = localPath + "/" + f.Name()
			}
			wg.Add(1)
			go downLoadFile(wg, sftpClient, f, localPath, fullRemotePath)
		}
		return
	}

	remoteFile, err := sftpClient.Open(remotePath)
	if err != nil {
		panic(err)
	}
	_, err = os.Stat(localPath)
	if err != nil {
		os.MkdirAll(localPath, 0766)
	}

	localFile, err := os.Create(fullLocalPath)
	if err != nil {
		panic(err)
	}
	defer localFile.Close()
	n, err := io.Copy(localFile, remoteFile)
	if err != nil {
		panic(err)
	}

	//获取远程文件大小
	log.Printf("文件下载成功[%s->%s]远程文件大小：%s，下载文件大小：%s",
		remotePath,
		fullLocalPath,
		formatFileSize(remoteFileInfo.Size()), formatFileSize(n))

}

func downLoadSingleFile(sftpClient *sftp.Client, f *DownloadFile) {

	remotePath := f.remotePath
	localPath := f.localPath
	remoteFile, err := sftpClient.Open(remotePath)
	if err != nil {
		panic(err)
	}
	runeStr := []rune(localPath)
	parentDir := string(runeStr[0:strings.LastIndex(localPath, "/")])
	_, err = os.Stat(parentDir)
	if err != nil {
		os.MkdirAll(parentDir, 0766)
	}

	localFile, err := os.Create(localPath)
	if err != nil {
		panic(err)
	}
	defer localFile.Close()
	n, err := io.Copy(localFile, remoteFile)
	if err != nil {
		panic(err)
	}

	//获取远程文件大小
	log.Printf("文件下载成功[%s->%s]远程文件大小：%s，下载文件大小：%s",
		remotePath,
		localPath,
		formatFileSize(f.size), formatFileSize(n))

}

func GetFiles(sftpClient *sftp.Client, localPath string, remotePath string) []*DownloadFile {
	remoteFileInfo, err := sftpClient.Stat(remotePath)
	var files []*DownloadFile
	if err != nil {
		panic(err)
	}
	//判断目录为空
	if remoteFileInfo.IsDir() {
		if dir, err := sftpClient.ReadDir(remotePath); err == nil {
			for _, f := range dir {
				files = append(files, GetFiles(sftpClient, localPath+"/"+f.Name(), remotePath+"/"+f.Name())...)
			}
		}
	} else {
		files = append(files, &DownloadFile{localPath: localPath, remotePath: remotePath, size: remoteFileInfo.Size()})
	}
	return files
}

func Upload(host string, port string, user string, pwd string) {
	sshConfig := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(pwd),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		ClientVersion:   "",
		Timeout:         10 * time.Second,
	}

	//建立与SSH服务器的连接
	sshClient, err := ssh.Dial("tcp", host+":"+port, sshConfig)
	if err != nil {
		log.Fatalln(err.Error())
	}
	defer sshClient.Close()

	sftpClient, err := sftp.NewClient(sshClient)
	if err != nil {
		log.Fatalln(err.Error())
	}

	defer sftpClient.Close()

	//获取当前目录
	cwd, err := sftpClient.Getwd()
	if err != nil {
		log.Fatalln(err.Error())
	}
	log.Println("当前目录：", cwd)

	//显示文件/目录详情
	fi, err := sftpClient.Lstat(cwd)
	log.Println(fi)

	{
		//上传文件(将本地file.dat文件通过sftp传到远程服务器)
		remoteFileName := fmt.Sprintf("upload_by_sftp_%d.dat", rand.Int())
		remoteFile, err := sftpClient.Create(sftp.Join(cwd, remoteFileName))
		if err != nil {
			log.Fatalln(err.Error())
		}
		defer remoteFile.Close()

		localFileName := "file.dat"
		//打开本地文件file.dat
		localFile, err := os.Open(localFileName)
		if err != nil {
			log.Fatalln(err.Error())
		}
		defer localFile.Close()

		//本地文件流拷贝到上传文件流
		n, err := io.Copy(remoteFile, localFile)
		if err != nil {
			log.Fatalln(err.Error())
		}

		//获取本地文件大小
		localFileInfo, err := os.Stat(localFileName)
		if err != nil {
			log.Fatalln(err.Error())
		}

		log.Printf("文件上传成功[%s->%s]本地文件大小：%s，上传文件大小：%s", localFileName, remoteFileName, formatFileSize(localFileInfo.Size()), formatFileSize(n))

		//计算文件MD5
		//windows计算：certutil -hashfile .\file.dat MD5
		//linux计算：md5sum upload_by_sftp_5577006791947779410.dat
	}
}

// 字节的单位转换 保留两位小数
func formatFileSize(s int64) (size string) {
	if s < 1024 {
		return fmt.Sprintf("%.2fB", float64(s)/float64(1))
	} else if s < (1024 * 1024) {
		return fmt.Sprintf("%.2fKB", float64(s)/float64(1024))
	} else if s < (1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fMB", float64(s)/float64(1024*1024))
	} else if s < (1024 * 1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fGB", float64(s)/float64(1024*1024*1024))
	} else if s < (1024 * 1024 * 1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fTB", float64(s)/float64(1024*1024*1024*1024))
	} else { //if s < (1024 * 1024 * 1024 * 1024 * 1024 * 1024)
		return fmt.Sprintf("%.2fEB", float64(s)/float64(1024*1024*1024*1024*1024))
	}
}
